//
// Created by syshen on 2021/4/17.
//

#ifndef FACE_DETECTOR_H
#define FACE_DETECTOR_H

#include <opencv2/opencv.hpp>
#include <string>
#include <stack>
#include <chrono>

#include "MNN/Interpreter.hpp"
#include "base_util.h"

#include "MNN/MNNDefine.h"
#include "MNN/Tensor.hpp"
#include "MNN/ImageProcess.hpp"

using namespace std::chrono;

//typedef struct location {
//    float _x;
//    float _y;
//}location;

//typedef struct landmark {
//    float x1;
//    float y1;
//    float x2;
//    float y2;
//    float s;
//    location point[5];
//}landmark;

//typedef struct rect_box {
//    float cx;
//    float cy;
//    float sx;
//    float sy;
//}rect_box;

class Retinaface
{

public:
	Retinaface() = delete;

	Retinaface(
		std::string mnn_path, int img_size, float threshold = 0.6,
		float nms = 0.4, int num_thread = 2, bool retinaface = false);

	int face_detect(cv::Mat& bgr, std::vector<landmark>& boxes);

	int face_detect(cv::Mat& img, landmark &face);

	~Retinaface();
private:

	void create_anchor(std::vector<rect_box> &anchor, int w, int h);

	void create_anchor_retinaface(std::vector<rect_box> &anchor, int w, int h);

	void compose_results(
		cv::Mat img, MNN::Tensor *scores, MNN::Tensor *boxes, MNN::Tensor *landmarks, std::vector<landmark> &results);

public:
	float _nms;
	float _threshold;
	float _mean_val[3];
	float _norm_vals[3] = { 1.0, 1.0, 1.0 };
	int _num_thread;
	int img_size;
	bool _retinaface;

	std::shared_ptr<MNN::Interpreter> det_interpreter;
	MNN::Session *det_session = nullptr;
	MNN::Tensor *input_tensor = nullptr;
	std::shared_ptr<MNN::CV::ImageProcess> pretreat = nullptr;
	std::vector<rect_box> anchors;
};

class FaceLandMark
{

public:
    FaceLandMark() = delete;

    FaceLandMark(std::string mnn_path, int img_size, int num_thread);

    void landmark_detect(cv::Mat& bgr, cv::Rect &rect, std::vector<cv::Point>& left_corner, \
                         std::vector<cv::Point>& right_corner);

    ~FaceLandMark();
private:
    std::shared_ptr<MNN::Interpreter> landmark_interpreter;
    MNN::Session *landmark_session = nullptr;
    MNN::Tensor *input_tensor = nullptr;
    std::shared_ptr<MNN::CV::ImageProcess> pretreat = nullptr;

    int img_size;
    int _num_thread;
    float _mean_val[3] = {0.0f, 0.0f, 0.0f};
    float _norm_vals[3] = { 1.0f, 1.0f, 1.0f};
};

class EyeSegmentation
{

public:
    EyeSegmentation() = delete;

    EyeSegmentation(std::string mnn_path, int img_size, int num_thread);

    cv::Mat eye_segment(const cv::Mat& bgr, const cv::Rect &rect, cv::Point &centroid);

    ~EyeSegmentation();
private:
    std::shared_ptr<MNN::Interpreter> interpreter;
    MNN::Session *session = nullptr;
    MNN::Tensor *input_tensor = nullptr;
    std::shared_ptr<MNN::CV::ImageProcess> pretreat = nullptr;

private:
    int img_size;
    int height, width;
    int _num_thread;
    float _mean_val[3] = {0.0f, 0.0f, 0.0f};
    float _norm_vals[3] = { 1.0f/255.0f, 1.0f/255.0f, 1.0f/255.0f};
//    float _norm_vals[3] = { 1.0f, 1.0f, 1.0f};
};

#endif //
